import argparse


def get_args():
    parser = argparse.ArgumentParser()
    # Required parameters
    parser.add_argument("--logger_name",
                        type=str,
                        default='root',
                        help="Name of this run. Used for monitoring.")
    parser.add_argument("--dataset",
                        choices=["cub", "dogs", "aircrafts"],
                        default="dogs",
                        help="which dataset to choose.")
    parser.add_argument('--data_root',
                        type=str,
                        # required=True,
                        help="where the dataset is")
    parser.add_argument("--batch_size",
                        type=int,
                        default=32)
    parser.add_argument("--epoch",
                        type=int,
                        default=30)
    parser.add_argument("--lr",
                        type=float)
    # parser.add_argument("--seed",
    #                     type=int,
    #                     default=100)
    # parser.add_argument("--n_gpu",
    #                     type=int,
    #                     default=1)
    parser.add_argument("--save_model_path",
                        type=str,
                        default='./model')
    parser.add_argument("--optim",
                        choices=['sgd', 'adam'],
                        type=str,
                        default='sgd')
    parser.add_argument("--loss_fn",
                        choices=['crossentropy', 'mse'],
                        type=str,
                        default='crossentropy')
    parser.add_argument("--pretrained_path",
                        type=str,
                        default='')
    parser.add_argument("--model",
                        choices=['resnet50', 'vit-b-16', 'vit-l-16', 'swin-t', 'swin-b'],
                        type=str,
                        default='vit-l-16')
    parser.add_argument("--lr_scheduler",
                        choices=['exp', 'cos'],
                        type=str,
                        default='')
    parser.add_argument("--image_size",
                        choices=[448, 224],
                        type=int,
                        default=448)

    args = parser.parse_args()

    return args
